from flask import request, render_template, jsonify,session,Blueprint
import requests
import json
import os
import time
from dotenv import load_dotenv

ai = Blueprint('ai', __name__)

# 加载环境变量
load_dotenv()

# 从环境变量获取API密钥
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
API_ENDPOINT = "https://api.deepseek.com/v1/chat/completions"

# 使用DeepSeek API生成文章
def generate_article_with_deepseek(topic, style, length):
    # 根据风格和长度准备不同的提示词
    style_prompts = {
        "academic": f"请以学术论文的风格，撰写一篇关于{topic}的文章。文章应该包含引言、正文和结论部分，使用客观、严谨的语言。",
        "blog": f"请以博客随笔的风格，撰写一篇关于{topic}的文章。语言应该轻松活泼，可以加入个人观点和经历。",
        "news": f"请以新闻报道的风格，撰写一篇关于{topic}的文章。文章应该客观公正，包含各方观点，使用第三人称。",
        "story": f"请以故事叙述的风格，创作一个关于{topic}的短篇故事。可以包含人物对话、情节发展和故事结局。"
    }
    
    length_desc = {
        "short": "文章长度约500字。",
        "medium": "文章长度约1000字。", 
        "long": "文章长度约2000字。"
    }
    
    # 组合提示词
    prompt = f"{style_prompts.get(style, style_prompts['blog'])} {length_desc.get(length, length_desc['medium'])} 请使用中文撰写，分段落呈现。"
    
    # 准备API请求
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}"
    }
    
    data = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": "你是一个专业的内容创作者，擅长根据主题创作各种风格的文章。"},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.7,
        "max_tokens": 2048
    }
    
    try:
        response = requests.post(API_ENDPOINT, headers=headers, data=json.dumps(data))
        response.raise_for_status()  # 检查请求是否成功
        
        result = response.json()
        content = result["choices"][0]["message"]["content"]
        return content
        
    except requests.exceptions.RequestException as e:
        raise Exception(f"API请求失败: {str(e)}")
    except (KeyError, IndexError) as e:
        raise Exception(f"解析API响应失败: {str(e)}")

# 处理AI生成请求
def handle_ai_generate():
    if request.method == 'POST':
        data = request.get_json()
        topic = data.get('topic', '')
        style = data.get('style', 'blog')
        length = data.get('length', 'medium')
        
        if not topic:
            return jsonify({'success': False, 'message': '主题不能为空'})
        
        try:
            content = generate_article_with_deepseek(topic, style, length)
            return jsonify({'success': True, 'content': content})
        except Exception as e:
            return jsonify({'success': False, 'message': str(e)})
    
    return jsonify({'success': False, 'message': '请求方法不支持'})

# 渲染AI写作页面
def show_ai_writing_page():
    return render_template('ai_writing.html')

#限流函数
   # 在AI控制器中添加请求计数器
   
   # 简单的内存计数器，实际项目中应使用Redis等持久化方案
request_counter = {}
   
def check_rate_limit(user_id):
    current_time = time.time()
       # 清理过期记录
    for uid in list(request_counter.keys()):
        if current_time - request_counter[uid]['timestamp'] > 3600:  # 1小时过期
            del request_counter[uid]
               
    if user_id not in request_counter:
        request_counter[user_id] = {'count': 1, 'timestamp': current_time}
        return True
       
    if request_counter[user_id]['count'] >= 10:  # 每小时限制10次请求
        return False
       
    request_counter[user_id]['count'] += 1
    return True




@ai.route('/ai_writing')
def ai_writing():
    return show_ai_writing_page()

@ai.route('/ai_generate', methods=['POST'])
def ai_generate():
    return handle_ai_generate()